import os
import requests
import sys
import csv
import colorama
from colorama import Fore, Style
from bs4 import BeautifulSoup
from export.export import export_to_excel
from lib.regex import is_domain


class RapidDnsScan:
    def __init__(self, domain: str):
        self.domain = domain
        self.url = "https://rapiddns.io/subdomain/{}?full=1#result".format(
            domain)
        self.soup = None
        self.results = []
        self.A = []
        self.MX = []
        self.AAAA = []
        self.subdomains = set()

    def request(self):
        logger.info(f'{self.domain} 子域名获取中...')
        page = requests.get(self.url)
        soup = BeautifulSoup(page.text, 'html.parser')
        self.soup = soup

    def parse(self):
        website_table = self.soup.find(
            "table", {"class": "table table-striped table-bordered"})
        if not website_table:
            return None
        tbody = website_table.tbody
        trs = tbody.find_all('tr')
        for tr in trs:
            tds = tr.find_all('td')
            domain = tds[0]
            address = tds[1]
            type = tds[2]
            result = {
                'subdomain': domain.string,
                'address': address.a.string,
                'type': type.string
            }
            if type.string == 'A':
                self.A.append(address.a.string)
            if type.string == 'AAAA':
                self.AAAA.append(address.a.string)
            if type.string == 'MX':
                self.MX.append(address.a.string)

            self.results.append(result)
            self.subdomains.add(domain.string)

    def export(self, save_path=None):
        columns_dict = {
            'subdomain': '子域名',
            'address': '解析结果',
            'type': '解析类型',
            'test': '测试'
        }
        export_to_excel(self.results, xlsx_name=self.domain, sheet_name='domain-' +
                        self.domain, columns_dict=columns_dict, save_path=save_path)

    def run(self, save_path):
        self.request()
        self.parse()
        self.export(save_path)


class RapidDNS:
    """FIXME: 批量域名查询, 当前需求不需要"""

    def __init__(self, domains: list):
        self.domains = domains
        self.ips = []
        self.subdomains = []

    def run(self, save_path):
        for domain in self.domains:
            if not is_domain(domain):
                continue
            rapid = RapidDnsScan(domain)
            rapid.request()
            rapid.parse()
            print(f'子域名数量: {len(rapid.results)}')
            if len(rapid.results) == 0:
                continue
            self.ips.extend(rapid.A)
            rapid.export(save_path)
